package com.amzing.dynamicrulesshardingjdbc.config.datasource;

import com.google.common.base.Preconditions;
import lombok.Generated;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;
import org.apache.shardingsphere.core.rule.ShardingRule;
import org.apache.shardingsphere.shardingjdbc.jdbc.adapter.AbstractDataSourceAdapter;
import org.apache.shardingsphere.shardingjdbc.jdbc.core.connection.ShardingConnection;
import org.apache.shardingsphere.shardingjdbc.jdbc.core.context.ShardingRuntimeContext;
import org.apache.shardingsphere.shardingjdbc.jdbc.core.datasource.MasterSlaveDataSource;
import org.apache.shardingsphere.shardingjdbc.jdbc.core.datasource.ShardingDataSource;
import org.apache.shardingsphere.spi.NewInstanceServiceLoader;
import org.apache.shardingsphere.transaction.core.TransactionTypeHolder;
import org.apache.shardingsphere.underlying.merge.engine.ResultProcessEngine;
import org.apache.shardingsphere.underlying.rewrite.context.SQLRewriteContextDecorator;
import org.apache.shardingsphere.underlying.route.decorator.RouteDecorator;

import javax.sql.DataSource;
import java.sql.SQLException;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicLong;
import java.util.stream.Collectors;

/**
 * 动态修改分表规则核心类
 *
 * @see ShardingDataSource
 */
@Slf4j
public class DynamicRulesDatasource extends AbstractDataSourceAdapter {

    private AtomicLong version;
    private Map<Long, ShardingRuntimeContext> shardingRuntimeContextMap;

    @Getter
    private Properties properties;

    public DynamicRulesDatasource(Map<String, DataSource> dataSourceMap, ShardingRule shardingRule, Properties props) throws SQLException {
        super(dataSourceMap);
        this.checkDataSourceType(dataSourceMap);
        ShardingRuntimeContext runtimeContext = new ShardingRuntimeContext(dataSourceMap, shardingRule, props, this.getDatabaseType());
        version = new AtomicLong(0);
        properties = props;
        shardingRuntimeContextMap = new ConcurrentHashMap<>();
        shardingRuntimeContextMap.put(version.getAndIncrement(), runtimeContext);

    }

    private void checkDataSourceType(Map<String, DataSource> dataSourceMap) {
        Iterator var2 = dataSourceMap.values().iterator();

        while (var2.hasNext()) {
            DataSource each = (DataSource) var2.next();
            Preconditions.checkArgument(!(each instanceof MasterSlaveDataSource), "Initialized data sources can not be master-slave data sources.");
        }

    }

    public final ShardingConnection getConnection() {

        return new ShardingConnection(this.getDataSourceMap(), getRuntimeContext(), TransactionTypeHolder.get());
    }

    /**
     * 设置分库分表规则
     * @param shardingRule
     * @param props
     * @throws SQLException
     */
    public void setRuntimeContext(ShardingRule shardingRule, Properties props) throws SQLException {
        ShardingRuntimeContext runtimeContext = new ShardingRuntimeContext(this.getDataSourceMap(), shardingRule, props, this.getDatabaseType());
        shardingRuntimeContextMap.put(version.getAndIncrement(),runtimeContext);
        this.resizeMap();
    }

    private void resizeMap(){

        Set<Long> keySet = shardingRuntimeContextMap.keySet();
        if(keySet.size() < 3){
            return;
        }

        List<Long> keys = keySet.stream().collect(Collectors.toList());
        keys.sort(Long::compareTo);

        //预留两个缓存，保证修改分表规则的时候，永远不会获取到null
        for (int i = 0; i < keys.size() - 2; i++) {
            long key = keys.get(i);
            if(!this.shardingRuntimeContextMap.containsKey(key)){
                continue;
            }
            this.shardingRuntimeContextMap.remove(key);
        }
    }

    @Generated
    public ShardingRuntimeContext getRuntimeContext() {
        //永远取最新的分表规则
        return shardingRuntimeContextMap.get(version.get() == 0 ? 0: version.get() - 1);
    }

    static {
        NewInstanceServiceLoader.register(RouteDecorator.class);
        NewInstanceServiceLoader.register(SQLRewriteContextDecorator.class);
        NewInstanceServiceLoader.register(ResultProcessEngine.class);
    }

}
